{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ctypes\n",
    "import llama_cpp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "llama_cpp.llama_backend_init(numa=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "llama_model_loader: loaded meta data with 20 key-value pairs and 291 tensors from /workspaces/llama-cpp-python/mistral-7b-v0.1.Q2_K.gguf (version GGUF V2)\n",
      "llama_model_loader: Dumping metadata keys/values. Note: KV overrides do not apply in this output.\n",
      "llama_model_loader: - kv   0:                       general.architecture str              = llama\n",
      "llama_model_loader: - kv   1:                               general.name str              = mistralai_mistral-7b-v0.1\n",
      "llama_model_loader: - kv   2:                       llama.context_length u32              = 32768\n",
      "llama_model_loader: - kv   3:                     llama.embedding_length u32              = 4096\n",
      "llama_model_loader: - kv   4:                          llama.block_count u32              = 32\n",
      "llama_model_loader: - kv   5:                  llama.feed_forward_length u32              = 14336\n",
      "llama_model_loader: - kv   6:                 llama.rope.dimension_count u32              = 128\n",
      "llama_model_loader: - kv   7:                 llama.attention.head_count u32              = 32\n",
      "llama_model_loader: - kv   8:              llama.attention.head_count_kv u32              = 8\n",
      "llama_model_loader: - kv   9:     llama.attention.layer_norm_rms_epsilon f32              = 0.000010\n",
      "llama_model_loader: - kv  10:                       llama.rope.freq_base f32              = 10000.000000\n",
      "llama_model_loader: - kv  11:                          general.file_type u32              = 10\n",
      "llama_model_loader: - kv  12:                       tokenizer.ggml.model str              = llama\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "llama_model_loader: - kv  13:                      tokenizer.ggml.tokens arr[str,32000]   = [\"<unk>\", \"<s>\", \"</s>\", \"<0x00>\", \"<...\n",
      "llama_model_loader: - kv  14:                      tokenizer.ggml.scores arr[f32,32000]   = [0.000000, 0.000000, 0.000000, 0.0000...\n",
      "llama_model_loader: - kv  15:                  tokenizer.ggml.token_type arr[i32,32000]   = [2, 3, 3, 6, 6, 6, 6, 6, 6, 6, 6, 6, ...\n",
      "llama_model_loader: - kv  16:                tokenizer.ggml.bos_token_id u32              = 1\n",
      "llama_model_loader: - kv  17:                tokenizer.ggml.eos_token_id u32              = 2\n",
      "llama_model_loader: - kv  18:            tokenizer.ggml.unknown_token_id u32              = 0\n",
      "llama_model_loader: - kv  19:               general.quantization_version u32              = 2\n",
      "llama_model_loader: - type  f32:   65 tensors\n",
      "llama_model_loader: - type q2_K:   65 tensors\n",
      "llama_model_loader: - type q3_K:  160 tensors\n",
      "llama_model_loader: - type q6_K:    1 tensors\n",
      "llm_load_vocab: special_eos_id is not in special_eog_ids - the tokenizer config may be incorrect\n",
      "llm_load_vocab: special tokens cache size = 3\n",
      "llm_load_vocab: token to piece cache size = 0.1637 MB\n",
      "llm_load_print_meta: format           = GGUF V2\n",
      "llm_load_print_meta: arch             = llama\n",
      "llm_load_print_meta: vocab type       = SPM\n",
      "llm_load_print_meta: n_vocab          = 32000\n",
      "llm_load_print_meta: n_merges         = 0\n",
      "llm_load_print_meta: vocab_only       = 0\n",
      "llm_load_print_meta: n_ctx_train      = 32768\n",
      "llm_load_print_meta: n_embd           = 4096\n",
      "llm_load_print_meta: n_layer          = 32\n",
      "llm_load_print_meta: n_head           = 32\n",
      "llm_load_print_meta: n_head_kv        = 8\n",
      "llm_load_print_meta: n_rot            = 128\n",
      "llm_load_print_meta: n_swa            = 0\n",
      "llm_load_print_meta: n_embd_head_k    = 128\n",
      "llm_load_print_meta: n_embd_head_v    = 128\n",
      "llm_load_print_meta: n_gqa            = 4\n",
      "llm_load_print_meta: n_embd_k_gqa     = 1024\n",
      "llm_load_print_meta: n_embd_v_gqa     = 1024\n",
      "llm_load_print_meta: f_norm_eps       = 0.0e+00\n",
      "llm_load_print_meta: f_norm_rms_eps   = 1.0e-05\n",
      "llm_load_print_meta: f_clamp_kqv      = 0.0e+00\n",
      "llm_load_print_meta: f_max_alibi_bias = 0.0e+00\n",
      "llm_load_print_meta: f_logit_scale    = 0.0e+00\n",
      "llm_load_print_meta: n_ff             = 14336\n",
      "llm_load_print_meta: n_expert         = 0\n",
      "llm_load_print_meta: n_expert_used    = 0\n",
      "llm_load_print_meta: causal attn      = 1\n",
      "llm_load_print_meta: pooling type     = 0\n",
      "llm_load_print_meta: rope type        = 0\n",
      "llm_load_print_meta: rope scaling     = linear\n",
      "llm_load_print_meta: freq_base_train  = 10000.0\n",
      "llm_load_print_meta: freq_scale_train = 1\n",
      "llm_load_print_meta: n_ctx_orig_yarn  = 32768\n",
      "llm_load_print_meta: rope_finetuned   = unknown\n",
      "llm_load_print_meta: ssm_d_conv       = 0\n",
      "llm_load_print_meta: ssm_d_inner      = 0\n",
      "llm_load_print_meta: ssm_d_state      = 0\n",
      "llm_load_print_meta: ssm_dt_rank      = 0\n",
      "llm_load_print_meta: ssm_dt_b_c_rms   = 0\n",
      "llm_load_print_meta: model type       = 7B\n",
      "llm_load_print_meta: model ftype      = Q2_K - Medium\n",
      "llm_load_print_meta: model params     = 7.24 B\n",
      "llm_load_print_meta: model size       = 2.87 GiB (3.41 BPW) \n",
      "llm_load_print_meta: general.name     = mistralai_mistral-7b-v0.1\n",
      "llm_load_print_meta: BOS token        = 1 '<s>'\n",
      "llm_load_print_meta: EOS token        = 2 '</s>'\n",
      "llm_load_print_meta: UNK token        = 0 '<unk>'\n",
      "llm_load_print_meta: LF token         = 13 '<0x0A>'\n",
      "llm_load_print_meta: EOG token        = 2 '</s>'\n",
      "llm_load_print_meta: max token length = 48\n",
      "llm_load_tensors: ggml ctx size =    0.14 MiB\n",
      "llm_load_tensors:        CPU buffer size =  2939.57 MiB\n",
      "..................................................................................................\n"
     ]
    }
   ],
   "source": [
    "params = llama_cpp.llama_model_default_params()\n",
    "params.n_gpu_layers = 35\n",
    "model = llama_cpp.llama_load_model_from_file(\n",
    "    b\"/workspaces/llama-cpp-python/mistral-7b-v0.1.Q2_K.gguf\", params\n",
    ")  # Update this to whatever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 415, 2936, 9060, 285, 1142]\n",
      "58\n"
     ]
    }
   ],
   "source": [
    "n_ctx = 512\n",
    "n_len = 32\n",
    "n_parallel = 2\n",
    "prompt = b\"The quick brown fox\"\n",
    "\n",
    "tokens = (llama_cpp.llama_token * n_ctx)()\n",
    "tokens_len = llama_cpp.llama_tokenize(\n",
    "    model, prompt, len(prompt), tokens, len(tokens), True, True\n",
    ")\n",
    "print(tokens[:tokens_len])\n",
    "\n",
    "n_kv_req = tokens_len + (n_len - tokens_len) * n_parallel\n",
    "print(n_kv_req)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "llama_new_context_with_model: n_ctx      = 64\n",
      "llama_new_context_with_model: n_batch    = 32\n",
      "llama_new_context_with_model: n_ubatch   = 32\n",
      "llama_new_context_with_model: flash_attn = 0\n",
      "llama_new_context_with_model: freq_base  = 10000.0\n",
      "llama_new_context_with_model: freq_scale = 1\n",
      "llama_kv_cache_init:        CPU KV buffer size =     8.00 MiB\n",
      "llama_new_context_with_model: KV self size  =    8.00 MiB, K (f16):    4.00 MiB, V (f16):    4.00 MiB\n",
      "llama_new_context_with_model:        CPU  output buffer size =     0.12 MiB\n",
      "llama_new_context_with_model:        CPU compute buffer size =     5.01 MiB\n",
      "llama_new_context_with_model: graph nodes  = 1030\n",
      "llama_new_context_with_model: graph splits = 1\n"
     ]
    }
   ],
   "source": [
    "ctx_params = llama_cpp.llama_context_default_params()\n",
    "ctx_params.seed = 1234\n",
    "ctx_params.n_ctx = n_kv_req\n",
    "ctx_params.n_batch = max(n_len, n_parallel)\n",
    "ctx_params.n_threads = 1\n",
    "ctx_params.n_threads_batch = 1\n",
    "ctx = llama_cpp.llama_new_context_with_model(model, ctx_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_ctx = llama_cpp.llama_n_ctx(ctx)\n",
    "batch = llama_cpp.llama_batch_init(max(tokens_len, n_parallel), 0, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "batch.n_tokens = tokens_len\n",
    "for i in range(tokens_len):\n",
    "    batch.token[i] = tokens[i]\n",
    "    batch.pos[i] = i\n",
    "    batch.seq_id[i][0] = 0\n",
    "    batch.n_seq_id[i] = 1\n",
    "    batch.logits[i] = False\n",
    "\n",
    "batch.logits[batch.n_tokens - 1] = True\n",
    "\n",
    "if llama_cpp.llama_decode(ctx, batch) != 0:\n",
    "    print(\"Error decoding\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(n_parallel):\n",
    "    llama_cpp.llama_kv_cache_seq_cp(ctx, 0, i, 0, batch.n_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Initialize sampler chain with default parameters\n",
    "sparams = llama_cpp.llama_sampler_chain_default_params()\n",
    "sampler_chain = llama_cpp.llama_sampler_chain_init(sparams)\n",
    "\n",
    "# Add top_k, top_p, temperature, and final distribution-based sampler\n",
    "llama_cpp.llama_sampler_chain_add(sampler_chain, llama_cpp.llama_sampler_init_top_k(40))\n",
    "llama_cpp.llama_sampler_chain_add(sampler_chain, llama_cpp.llama_sampler_init_top_p(0.9, 1))\n",
    "llama_cpp.llama_sampler_chain_add(sampler_chain, llama_cpp.llama_sampler_init_temp(0.4))\n",
    "llama_cpp.llama_sampler_chain_add(sampler_chain, llama_cpp.llama_sampler_init_dist(1234))  # Final \"dist\" sampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7\n",
      "[' j', ' jumped']\n",
      "8\n",
      "[' j over', ' jumped over']\n",
      "9\n",
      "[' j over the', ' jumped over the']\n",
      "10\n",
      "[' j over the lazy', ' jumped over the lazy']\n",
      "11\n",
      "[' j over the lazy dog', ' jumped over the lazy dog']\n",
      "12\n",
      "[' j over the lazy dog.', ' jumped over the lazy dog\\n']\n",
      "13\n",
      "[' j over the lazy dog. También', ' jumped over the lazy dog\\nGroupLayout']\n",
      "14\n",
      "[' j over the lazy dog. También:', ' jumped over the lazy dog\\nGroupLayouting']\n",
      "15\n",
      "[' j over the lazy dog. También: is', ' jumped over the lazy dog\\nGroupLayouting is']\n",
      "16\n",
      "[' j over the lazy dog. También: is a', ' jumped over the lazy dog\\nGroupLayouting is a']\n",
      "17\n",
      "[' j over the lazy dog. También: is a technique', ' jumped over the lazy dog\\nGroupLayouting is a common']\n",
      "18\n",
      "[' j over the lazy dog. También: is a technique practice', ' jumped over the lazy dog\\nGroupLayouting is a common practice']\n",
      "19\n",
      "[' j over the lazy dog. También: is a technique practice in', ' jumped over the lazy dog\\nGroupLayouting is a common practice in']\n",
      "20\n",
      "[' j over the lazy dog. También: is a technique practice in the', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the']\n",
      "21\n",
      "[' j over the lazy dog. También: is a technique practice in the real', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media']\n",
      "22\n",
      "[' j over the lazy dog. También: is a technique practice in the real-', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry']\n",
      "23\n",
      "[' j over the lazy dog. También: is a technique practice in the real-.', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry.']\n",
      "24\n",
      "[' j over the lazy dog. También: is a technique practice in the real-. We', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However']\n",
      "25\n",
      "[' j over the lazy dog. También: is a technique practice in the real-. We,', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However,']\n",
      "26\n",
      "[' j over the lazy dog. También: is a technique practice in the real-. We, when', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there']\n",
      "27\n",
      "[' j over the lazy dog. También: is a technique practice in the real-. We, when is', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there has']\n",
      "28\n",
      "[' j over the lazy dog. También: is a technique practice in the real-. We, when is been', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there has been']\n",
      "29\n",
      "[' j over the lazy dog. También: is a technique practice in the real-. We, when is been little', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there has been little']\n",
      "30\n",
      "[' j over the lazy dog. También: is a technique practice in the real-. We, when is been little research', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there has been little emp']\n",
      "31\n",
      "[' j over the lazy dog. También: is a technique practice in the real-. We, when is been little researchirical', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there has been little empirical']\n",
      "32\n",
      "[' j over the lazy dog. También: is a technique practice in the real-. We, when is been little researchirical research', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there has been little empirical research']\n"
     ]
    }
   ],
   "source": [
    "streams = [\"\"] * n_parallel\n",
    "i_batch = [batch.n_tokens - 1] * n_parallel\n",
    "\n",
    "n_cur = batch.n_tokens\n",
    "n_decode = 0\n",
    "\n",
    "while n_cur <= n_len:\n",
    "    batch.n_tokens = 0\n",
    "    for i in range(n_parallel):\n",
    "        if i_batch[i] < 0:\n",
    "            continue\n",
    "\n",
    "        # Sample the next token using the sampler chain\n",
    "        new_token_id = llama_cpp.llama_sampler_sample(sampler_chain, ctx, -1)\n",
    "\n",
    "        if new_token_id == llama_cpp.llama_token_eos(ctx) or n_cur == n_len:\n",
    "            i_batch[i] = -1\n",
    "            continue\n",
    "\n",
    "        buf = (ctypes.c_char * 32)()\n",
    "        \n",
    "        # Convert token ID to text\n",
    "        outlen = llama_cpp.llama_token_to_piece(model, new_token_id, buf, len(buf), 0, False)\n",
    "        streams[i] += bytes(buf[:outlen]).decode(\"utf-8\")\n",
    "\n",
    "        batch.token[batch.n_tokens] = new_token_id\n",
    "        batch.pos[batch.n_tokens] = n_cur\n",
    "        batch.seq_id[batch.n_tokens][0] = i\n",
    "        batch.n_seq_id[batch.n_tokens] = 1\n",
    "        batch.logits[batch.n_tokens] = True\n",
    "\n",
    "        i_batch[i] = batch.n_tokens\n",
    "        batch.n_tokens += 1\n",
    "        n_decode += 1\n",
    "\n",
    "    if batch.n_tokens == 0:\n",
    "        break\n",
    "\n",
    "    n_cur += 1\n",
    "\n",
    "    if llama_cpp.llama_decode(ctx, batch) != 0:\n",
    "        print(\"Error decoding\", flush=True)\n",
    "        break\n",
    "    print(n_cur)\n",
    "    print(streams)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[' j over the lazy dog. También: is a technique practice in the real-. We, when is been little researchirical research', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there has been little empirical research']\n"
     ]
    }
   ],
   "source": [
    "print(streams)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "llama_cpp.llama_batch_free(batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "llama_cpp.llama_free(ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "llama_cpp.llama_free_model(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "llama_cpp.llama_backend_free()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
